package com.sosky;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * @author huanghuanlono
 * @date 2021/12/6 18:22
 * <p>
 * 1424. 对角线遍历 II
 * <p>
 * 给你一个列表 nums ，里面每一个元素都是一个整数列表。请你依照下面各图的规则，按顺序返回 nums 中对角线上的整数。
 * <p>
 * 示例 1：
 * 输入：nums = [[1,2,3],[4,5,6],[7,8,9]]
 * 输出：[1,4,2,7,5,3,8,6,9]
 * <p>
 * 示例 2：
 * 输入：nums = [[1,2,3,4,5],[6,7],[8],[9,10,11],[12,13,14,15,16]]
 * 输出：[1,6,2,8,7,3,9,4,12,10,5,13,11,14,15,16]
 * <p>
 * 示例 3：
 * 输入：nums = [[1,2,3],[4],[5,6,7],[8],[9,10,11]]
 * 输出：[1,4,2,5,3,8,6,9,7,10,11]
 * <p>
 * 示例 4：
 * 输入：nums = [[1,2,3,4,5,6]]
 * 输出：[1,2,3,4,5,6]
 * <p>
 * 提示：
 * 1 <= nums.length <= 10^5
 * 1 <= nums[i].length <=10^5
 * 1 <= nums[i][j] <= 10^9
 * nums中最多有10^5个数字。
 */
public class Solution1424 {
    public static void main(String[] args) {
        List<List<Integer>> data1 = new ArrayList<>();
        data1.add(Arrays.asList(1, 2, 3));
        data1.add(Arrays.asList(4, 5, 6));
        data1.add(Arrays.asList(7, 8, 9));
//        System.out.println(Arrays.toString(findDiagonalOrder(data1)));
        List<List<Integer>> data2 = new ArrayList<>();
        data2.add(Arrays.asList(1, 2, 3, 4, 5));
        data2.add(Arrays.asList(6, 7));
        data2.add(Arrays.asList(8));
        data2.add(Arrays.asList(9, 10, 11));
        data2.add(Arrays.asList(12, 13, 14, 15, 16));
        System.out.println(Arrays.toString(findDiagonalOrder(data2)));
    }

    public static int[] findDiagonalOrder(List<List<Integer>> nums) {
        int totalCount = 0;
        int maxCapacity = 0;
        for (int i = 0; i < nums.size(); i++) {
            totalCount += nums.get(i).size();
            if (nums.get(i).size() > maxCapacity) {
                maxCapacity = nums.get(i).size();
            }
        }
        int[] result = new int[totalCount];
        int index = 0;
        int xIndexOfStart = 0;
        int yIndexOfStart = 0;
        int xIndex = 0;
        int yIndex = 0;
        while (true) {
            List<Integer> data = nums.get(xIndex);
            if (data != null) {
                if (yIndex < data.size()) {
                    //有值加入结果列表
                    result[index] = data.get(yIndex);
                    index++;
                    //遍历完毕跳出循环
                    if (index == result.length) {
                        break;
                    }
                }
                if (xIndex == 0 && yIndex < maxCapacity - 1) {
                    xIndexOfStart += 1;
                    xIndex = xIndexOfStart;
                    yIndex = 0;
                    yIndexOfStart = yIndex;
                } else if (yIndex >= maxCapacity - 1) {
                    if (xIndexOfStart == nums.size() - 1) {
                        xIndex = xIndexOfStart;
                    }
                    yIndex = yIndexOfStart + 1;
                    yIndexOfStart = yIndex;
                } else {
                    xIndex -= 1;
                    yIndex += 1;
                }
            }
        }
        return result;
    }
}
